
#include "CPUAnalyzer.h"
#include "Common/Strings.h"
#include "DeviceAnalyzer.h"
#include <comutil.h>
#include <psapi.h>

using namespace std;
using namespace ttg;
using namespace ttg::services;

//stolen from winternl.h (to avoid requirement for DDK)

typedef LONG NTSTATUS;

typedef enum _SYSTEM_INFORMATION_CLASS {
    SystemBasicInformation = 0,
    SystemPerformanceInformation = 2,
    SystemTimeOfDayInformation = 3,
    SystemProcessInformation = 5,
    SystemProcessorPerformanceInformation = 8,
    SystemInterruptInformation = 23,
    SystemExceptionInformation = 33,
    SystemRegistryQuotaInformation = 37,
    SystemLookasideInformation = 45
} SYSTEM_INFORMATION_CLASS;

NTSTATUS (WINAPI *NtQuerySystemInformation) (
    IN SYSTEM_INFORMATION_CLASS SystemInformationClass,
    OUT PVOID SystemInformation,
    IN ULONG SystemInformationLength,
    OUT PULONG ReturnLength OPTIONAL
    );

//-------------------
//--- CPUAnalyzer ---
//-------------------

CPUAnalyzer *CPUAnalyzer::ref = new CPUAnalyzer();

std::string CPUAnalyzer::determineCPUDescription()
{
	HKEY rKey;
	if (RegOpenKeyExA(HKEY_LOCAL_MACHINE,
					  "HARDWARE\\DESCRIPTION\\System\\CentralProcessor\\0",
					  0,
					  KEY_QUERY_VALUE,
					  &rKey) != ERROR_SUCCESS)
		return string("CPU not found.");
	else
	{
		DWORD size = 256;
		string res(size, '\0');
		if (RegQueryValueExA(rKey,
							 "ProcessorNameString",
							 NULL,
							 NULL,
							 (LPBYTE)res.c_str(),
							 &size) != ERROR_SUCCESS)
		{
			return string("CPU not found.");
		}
		RegCloseKey(rKey);
		return res;
	}
}

size_t CPUAnalyzer::determineCoreCount()
{
	SYSTEM_INFO sinfo;
	GetSystemInfo(&sinfo);
	return (size_t)sinfo.dwNumberOfProcessors;
}

size_t CPUAnalyzer::determineTotalPhysicalMemory()
{
	PERFORMANCE_INFORMATION pinfo;
	GetPerformanceInfo(&pinfo, sizeof(PERFORMANCE_INFORMATION));
	return (size_t)(pinfo.PhysicalTotal / (1024 * 1024.0f) * pinfo.PageSize);
}

size_t CPUAnalyzer::determineAvailablePhysicalMemory()
{
	PERFORMANCE_INFORMATION pinfo;
	GetPerformanceInfo(&pinfo, sizeof(PERFORMANCE_INFORMATION));
	return (size_t)(pinfo.PhysicalAvailable / (1024 * 1024.0f) * pinfo.PageSize);
}

void CPUAnalyzer::startUnSafe(std::vector<Service *> &deps, const char *params)
{
	//Getting DllLoader.
	if (deps.size() < 1 ||
		(this->loader = dynamic_cast<DllLoader *>(deps[0])) == NULL)
		throw Error(Error::FATAL, "CPUAnalyzer", "startUnSafe",
					"Cannot get DllLoader");
	loader->load("ntdll", "NtQuerySystemInformation", (FARPROC &)NtQuerySystemInformation);
	//Setting current usage.
	size_t count = getCoreCount();
	previous = new SYSTEM_PROCESSOR_PERFORMANCE_INFORMATION[count + 1];
	current = new SYSTEM_PROCESSOR_PERFORMANCE_INFORMATION[count + 1];
	ULONG req;
	if (NtQuerySystemInformation(SystemProcessorPerformanceInformation,
								 previous,
								 (ULONG)(sizeof(SYSTEM_PROCESSOR_PERFORMANCE_INFORMATION) * count),
								 &req) != 0)
		throw Error(Error::FATAL, "CPUAnalyzer", "CPUAnalyzer", "Cannot get info about CPU usage");
	t = 0.0f;
	//Resizing array with descriptions.
	core_usage.resize(count + 1);
	//Adding info to SettingKeeper
	Settings *s = new Settings("COUNT_OF_CORES", "Dynamic attribute.");
	s->add(s->setChosen(new SettingsUnit("", "", (int)count)));
	SettingsKeeper::getRef()->add(s);
	//Registering.
	DeviceAnalyzer::getRef()->registerAnalyzer(getDeviceInfo());
	NotifierID
		= Timer::getRef()->setNotifier(ttg::structures::make_delegate(this, &CPUAnalyzer::onTick), 0.5, true);
}

void CPUAnalyzer::stopUnSafe()
{
	Timer::getRef()->disableNotifier(NotifierID);
	DeviceAnalyzer::getRef()->unRegisterAnalyzer(getDeviceInfo());
	SAFE_DELETE_ARRAY(current);
	SAFE_DELETE_ARRAY(previous);
}

void CPUAnalyzer::onTick(double step)
{
	t += step;
	if (t > 0.5f)
	{
		t = 0;
		size_t count = getCoreCount();
		//Getting current times.
		ULONG req;
		if (NtQuerySystemInformation(SystemProcessorPerformanceInformation,
									 current,
									 (ULONG)(sizeof(SYSTEM_PROCESSOR_PERFORMANCE_INFORMATION) * count),
									 &req) != 0)
			throw Error(Error::FATAL, "CPUAnalyzer", "CPUAnalyzer", "Cannot get info about CPU usage");
		//Updating usage for each core and total usage.
		float summed_usage = 0.0f;
		for (size_t i = 0; i < count; i++)
		{
			LARGE_INTEGER user;
			user.QuadPart= current[i].UserTime.QuadPart - previous[i].UserTime.QuadPart;
			LARGE_INTEGER kernel;
			kernel.QuadPart = current[i].KernelTime.QuadPart - previous[i].KernelTime.QuadPart;
			LARGE_INTEGER idle;
			idle.QuadPart = current[i].IdleTime.QuadPart - previous[i].IdleTime.QuadPart;
			float usage = (float)(static_cast<double>(user.QuadPart + kernel.QuadPart - idle.QuadPart) /
								  static_cast<double>(user.QuadPart + kernel.QuadPart));
			info.getSubDevices()[i]->setUsage(usage);
			summed_usage += usage;
		}
		info.setUsage(summed_usage / getCoreCount());
		//Updating "memory" subdevice.
		SubDeviceInfo *memory_info = info.getSubDevices()[info.getSubDevices().size() - 1];
		memory_info->setUsage( 1.0f - (float)determineAvailablePhysicalMemory() / memory_info->getMaxValue());
		//Swapping times.
		SYSTEM_PROCESSOR_PERFORMANCE_INFORMATION *temp = previous;
		previous = current;
		current = temp;
		//Generating legends.
		this->setCoreUsages();
		this->setFullInfo();
	}
}

void CPUAnalyzer::setCoreUsages()
{
	for (size_t i = 0; i < core_usage.size(); i++)
	{
		char buf[256];
		float usage;
		if (i < getCoreCount())
		{
			usage = info.getSubDevices()[i]->getUsage();
			sprintf_s(buf, 256, "core #%2d:            %3d%%", i, static_cast<int>(usage * 100));
		}
		else
		{
			usage = info.getUsage();
			sprintf_s(buf, 256, "TOTAL :            %3d%%", static_cast<int>(usage * 100));
		}
		for (size_t j = 0; j < 10; j++)
			if (i < (usage + 0.005f) * 10 - 1.0f)
				buf[11 + j] = '#';
			else
				buf[11 + j] = '_';
		core_usage[i] = buf;
	}
}

void CPUAnalyzer::setFullInfo()
{
	char buf[256];
	sprintf_s(buf, 256, "%s\n", getCPUInfo());
	full_info = full_info + buf;
	for (size_t i = 0; i < core_usage.size(); i++)
	{
		sprintf_s(buf, 256, "%s\n", core_usage[i].c_str());
		full_info = full_info + buf;
	}
}

CPUAnalyzer::CPUAnalyzer()
	:Service("CPUAnalyzer", "DllLoader SettingsKeeper Timer"),
	info("CPU", determineCPUDescription().c_str(), "%", 100.0f)
{
	//Creating info for CPU cores.
	core_count = determineCoreCount();
	for (size_t i = 0; i < core_count; i++)
	{
		char buf[256];
		sprintf_s(buf, 256, "Core #%d", i);
		info.addSubDevice(new SubDeviceInfo(buf, "%", 100.0f));
	}
	//Creating info for memory.
	info.addSubDevice(new SubDeviceInfo("Physical memory", "MB", (float)determineTotalPhysicalMemory()));
	makeRegistration();
}

CPUAnalyzer::~CPUAnalyzer()
{
	this->stop();
	ref = NULL;
}
